import math
import time
import random
import re
import json
import requests
import pandas as pd


EXCHANGE_DICT = {"SSE": "m:10", "SZSE": "m:12"}

def option_risk_analysis_em(exchange: str = "SSE") -> pd.DataFrame:
    """
    ak的option/option_risk_analysis_em 用不了，只能查询sse的数据
    这里做了改造
    东方财富网-数据中心-特色数据-期权风险分析
    https://data.eastmoney.com/other/riskanal.html
    :return: 期权风险分析
    :rtype: pandas.DataFrame
    """
    if exchange not in EXCHANGE_DICT:
        raise ValueError("exchange 只能是 'SSE' 或 'SZSE'")

    fs = EXCHANGE_DICT[exchange]
    pz = 50
    base_url = (
        "https://push2.eastmoney.com/api/qt/clist/get?"
        "fid=f12&po=1&np=1&fltt=2&invt=2"
        "&ut=8dec03ba335b81bf4ebdf7b29ec27d15"
        "&fields=f1,f2,f3,f12,f13,f14,f302,f303,f325,f326,f327,f329,f328,f301,f152,f154"
        f"&fs={fs}&pz={pz}&pn={{page}}"
    )
    headers = {
        "Referer": "https://data.eastmoney.com/other/valueAnal.html",
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) "
            "Chrome/137.0.0.0 Safari/537.36"
        ),
    }

    session = requests.Session()
    resp = session.get(base_url.format(page=1), headers=headers, timeout=30)
    match = re.search(r"\{.*\}\}$", resp.text)
    if not match:
        raise ValueError("未找到有效 JSON 数据")
    js = json.loads(match.group())
    total = js["data"]["total"]
    pages = math.ceil(total / pz)

    data = js["data"]["diff"]

    for page in range(2, pages + 1):
        resp = session.get(base_url.format(page=page), headers=headers, timeout=30)
        match = re.search(r"\{.*\}\}$", resp.text)
        if not match:
            continue
        js = json.loads(match.group())
        data.extend(js["data"]["diff"])
        time.sleep(0.3 + random.random() * 0.2)
    temp_df = pd.DataFrame(data)

    # pd.set_option('display.max_columns', None)  # 显示所有列
    # temp_df.to_csv("tmp.csv", index=False, encoding='utf-8-sig')
    # print(temp_df)

    temp_df.columns = [
        "-",
        "最新价",
        "涨跌幅",
        "期权代码",
        "-",
        "期权名称",
        "-",
        "-",
        "到期日",
        "杠杆比率",
        "实际杠杆比率",
        "Delta",
        "Gamma",
        "Vega",
        "Theta",
        "Rho",
    ]
    temp_df = temp_df[
        [
            "期权代码",
            "期权名称",
            "最新价",
            "涨跌幅",
            "杠杆比率",
            "实际杠杆比率",
            "Delta",
            "Gamma",
            "Vega",
            "Rho",
            "Theta",
            "到期日",
        ]
    ]
    temp_df["最新价"] = round(pd.to_numeric(temp_df["最新价"], errors="coerce"), 4)
    temp_df["涨跌幅"] = pd.to_numeric(temp_df["涨跌幅"], errors="coerce")
    temp_df["杠杆比率"] = pd.to_numeric(temp_df["杠杆比率"], errors="coerce")
    temp_df["实际杠杆比率"] = pd.to_numeric(temp_df["实际杠杆比率"], errors="coerce")
    temp_df["Delta"] = pd.to_numeric(temp_df["Delta"], errors="coerce")
    temp_df["Gamma"] = pd.to_numeric(temp_df["Gamma"], errors="coerce")
    temp_df["Vega"] = pd.to_numeric(temp_df["Vega"], errors="coerce")
    temp_df["Rho"] = pd.to_numeric(temp_df["Rho"], errors="coerce")
    temp_df["Theta"] = pd.to_numeric(temp_df["Theta"], errors="coerce")
    temp_df["到期日"] = pd.to_datetime(
        temp_df["到期日"], format="%Y%m%d", errors="coerce"
    ).dt.date
    return temp_df

def fetch_all_option_risk() :
    print(f"start fetch risk...")
    df_sse = option_risk_analysis_em("SSE")
    print(f"SSE end")
    # print(df_sse)
    time.sleep(3)  # 防止请求过快
    df_szse = option_risk_analysis_em("SZSE")
    print(f"SZSE end")
    # print(df_szse)
    # 纵向合并两个DataFrame
    combined_df = pd.concat([df_sse, df_szse], ignore_index=True)
    return combined_df

if __name__ == "__main__":
    filename = 'option_risk_analysis_em.csv'
    combined_df = fetch_all_option_risk()
    combined_df.to_csv(filename, index=False, encoding='utf-8-sig')

    # combined_df = pd.read_csv(filename)
    print(combined_df)
